#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@Project : light-edit
@File : Assembling.py
@Author : Yj
@Time : 2025/9/9
@脚本说明 : 多功能扫描工具，整合端口扫描、MySQL、Redis、SSH、IP段和子域名扫描
"""
import argparse
import time
import ipaddress
from concurrent.futures import ThreadPoolExecutor
from Tools.port_scan import port_scanner
from Tools.mysql_scan import mysql_run
from Tools.redis_scan import redis_run
from Tools.ssh_scan import ssh_run
from Tools.ip_scan import scan_ip
from Tools.subdomain_scan import subdomain_run
from termcolor import colored

banner = """
 /$$      /$$           /$$   /$$     /$$        /$$$$$$                                                             
| $$$    /$$$          | $$  | $$    |__/       /$$__  $$                                                            	
| $$$$  /$$$$ /$$   /$$| $$ /$$$$$$   /$$      | $$  \__/  /$$$$$$$  /$$$$$$  /$$$$$$$  /$$$$$$$   /$$$$$$   /$$$$$$ 
| $$ $$/$$ $$| $$  | $$| $$|_  $$_/  | $$      |  $$$$$$  /$$_____/ |____  $$| $$__  $$| $$__  $$ /$$__  $$ /$$__  $$	
| $$  $$$| $$| $$  | $$| $$  | $$    | $$       \____  $$| $$        /$$$$$$$| $$  \ $$| $$  \ $$| $$$$$$$$| $$  \__/	
| $$\  $ | $$| $$  | $$| $$  | $$ /$$| $$       /$$  \ $$| $$       /$$__  $$| $$  | $$| $$  | $$| $$_____/| $$      
| $$ \/  | $$|  $$$$$$/| $$  |  $$$$/| $$      |  $$$$$$/|  $$$$$$$|  $$$$$$$| $$  | $$| $$  | $$|  $$$$$$$| $$      
|__/     |__/ \______/ |__/   \___/  |__//$$$$$$\______/  \_______/ \_______/|__/  |__/|__/  |__/ \_______/|__/      
                                        |______/                                                                
"""

print(colored(banner, 'magenta'))


def run_port(args):
    """运行端口扫描"""
    try:
        host = args.host
        start_port, end_port = map(int, args.port_range.split('-'))
        if not (1 <= start_port <= end_port <= 65535):
            raise ValueError("端口范围必须在 1-65535 之间！")
        start_time = time.time()
        port_scanner(host, start_port, end_port)
        print(f"扫描完成，耗时 {time.time() - start_time:.2f} 秒")
    except ValueError as e:
        print(f"输入错误：{e}")
    except Exception as e:
        print(f"发生错误：{e}")


def run_ip(args):
    """运行 IP 段扫描"""
    try:
        network = ipaddress.ip_network(args.host, strict=False)
        subnets = list(network.subnets(new_prefix=24)) if network.prefixlen < 24 else [network]
        ip_segments = [str(subnet) for subnet in subnets]
        if not ip_segments:
            print("IP段无效")
            return
        alive_count = 0
        start_time = time.time()
        with ThreadPoolExecutor(max_workers=args.max_workers) as executor:
            results = executor.map(scan_ip, ip_segments)
            for alive_hosts in results:
                alive_count += len(alive_hosts)
        print(f"[+] 存活主机数：{alive_count}")
        print(f"扫描完成，耗时 {int(time.time() - start_time)} 秒")
    except Exception as e:
        print(f"发生错误：{e}")


def run_mysql(host, port, users, passwords, max_workers=50):
    """运行 MySQL 扫描，支持单 IP 或网段"""
    try:
        start_time = time.time()
        if '/' in host:
            # 网段扫描
            network = ipaddress.ip_network(host, strict=False)
            with ThreadPoolExecutor(max_workers=max_workers) as executor:
                for ip in network.hosts():
                    ip_str = str(ip)
                    executor.submit(mysql_run, ip_str, port, users, passwords)
        else:
            # 单 IP 扫描
            mysql_run(host, port, users, passwords)
        print(f"MySQL 扫描完成，耗时 {int(time.time() - start_time)} 秒")
    except Exception as e:
        print(f"发生错误：{e}")


def main():
    parser = argparse.ArgumentParser(description="Multi_Scanner")
    subparsers = parser.add_subparsers(dest="sub_command", help="选择要使用的功能")

    # Port扫描解析器
    port_parser = subparsers.add_parser("port", help="端口扫描")
    port_parser.add_argument("--host", "-H", type=str, default="127.0.0.1", help="目标主机地址")
    port_parser.add_argument("--port-range", "-P", type=str, default="1-1000", help="端口范围（如 1-1000）")

    # MySQL 扫描解析器
    mysql_parser = subparsers.add_parser("mysql", help="MySQL弱口令扫描")
    mysql_parser.add_argument("--host", "-H", type=str, default="127.0.0.1",
                              help="MySQL主机地址或网段（如 192.168.1.0/24）")
    mysql_parser.add_argument("--port", "-P", type=int, default=3306, help="MySQL端口号")
    mysql_parser.add_argument("--username", "-u", type=argparse.FileType(mode='r', encoding="utf-8"),
                              default="./dict/user.txt", help="MySQL用户名字典文件")
    mysql_parser.add_argument("--password", "-p", type=argparse.FileType(mode='r', encoding="utf-8"),
                              default="./dict/passwd.txt", help="MySQL密码字典文件")
    mysql_parser.add_argument("--max-workers", type=int, default=50, help="最大线程数（网段扫描时使用）")

    # Redis 扫描解析器
    redis_parser = subparsers.add_parser("redis", help="Redis弱口令及未授权访问扫描")
    redis_parser.add_argument("--host", "-H", type=str, default="127.0.0.1", help="Redis主机地址")
    redis_parser.add_argument("--port", "-P", type=int, default=6379, help="Redis端口号")
    redis_parser.add_argument("--password", "-p", type=argparse.FileType(mode='r', encoding="utf-8"),
                              default="./dict/passwd.txt", help="Redis密码字典文件")

    # SSH 扫描解析器
    ssh_parser = subparsers.add_parser("ssh", help="SSH弱口令扫描")
    ssh_parser.add_argument("--ip", type=str, default="127.0.0.1", help="ip address")
    ssh_parser.add_argument("--port", "-P", type=int, default=22, help="SSH端口号")
    ssh_parser.add_argument("--username", "-u", type=argparse.FileType(mode='r', encoding="utf-8"),
                            default="./dict/user.txt", help="SSH用户名字典文件")
    ssh_parser.add_argument("--password", "-p", type=argparse.FileType(mode='r', encoding="utf-8"),
                            default="./dict/passwd.txt", help="SSH密码字典文件")

    # IP 段扫描解析器
    ip_parser = subparsers.add_parser("ip", help="IP段存活主机扫描")
    ip_parser.add_argument("--host", "-H", type=str, default="127.0.0.1/24", help="IP段（如 192.168.1.0/24）")
    ip_parser.add_argument("--max-workers", type=int, default=50, help="最大线程数")

    # 子域名扫描解析器
    subdomain_parser = subparsers.add_parser("subdomain", help="子域名扫描")
    subdomain_parser.add_argument("--domain", "-d", type=str, required=True, help="主域名（如 example.com）")
    subdomain_parser.add_argument("--subdomain-file", "-f", type=str, default="./dict/subdomain.txt", help="子域名字典文件")

    args = parser.parse_args()

    if not args.sub_command:
        parser.print_help()
        return

    if args.sub_command == "port":
        run_port(args)
    elif args.sub_command == "mysql":
        run_mysql(args.host, args.port,
                  [line.strip() for line in args.username.readlines()],
                  [line.strip() for line in args.password.readlines()],
                  args.max_workers)
    elif args.sub_command == "redis":
        redis_run(args.host, args.port, [line.strip() for line in args.password.readlines()])
    elif args.sub_command == "ssh":
        ssh_run(args.ip_list, args.port, args.username.name, args.password.name)
    elif args.sub_command == "ip":
        run_ip(args)
    elif args.sub_command == "subdomain":
        subdomain_run(args.domain, args.subdomain_file)
    else:
        print("无效的功能选择")


if __name__ == "__main__":
    main()
